{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0780ab9b-f47f-4251-9215-949fd1e5600c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# following https://www.datacamp.com/tutorial/chromadb-tutorial-step-by-step-guide"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c9cf7c62-5385-4bd7-9d06-19be62a119cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import chromadb\n",
    "from chromadb.config import Settings\n",
    "\n",
    "client = chromadb.PersistentClient(path=\"/home/jupyter/chromadb/data\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c1acf397-8365-45d7-b333-bedc1aa9da08",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 指定创建collection使用的embedding模型和dimension，保证collection.quy\n",
    "import os\n",
    "from chromadb.utils import embedding_functions\n",
    "collection = client.get_or_create_collection(\n",
    "    name=\"Employers\",\n",
    "    embedding_function=embedding_functions.OpenAIEmbeddingFunction(\n",
    "                model_name=\"text-embedding-v4\",\n",
    "                api_key=os.getenv(\"DASHSCOPE_API_KEY\"),\n",
    "                api_base=\"https://dashscope.aliyuncs.com/compatible-mode/v1\",\n",
    "                dimensions=1024,\n",
    "            )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "bab9c4a6-cc54-4ac1-8b0e-da0a0217be92",
   "metadata": {},
   "outputs": [],
   "source": [
    "employer1_info = \"\"\"\n",
    "我叫shinerio，来自北京，负责小米生态产品营销工作。\n",
    "\"\"\"\n",
    "\n",
    "employer2_info = \"\"\"\n",
    "Tom是华为公司终端的一名销售，负责问界汽车的营销工作。\n",
    "\"\"\"\n",
    "\n",
    "employer3_info = \"\"\"\n",
    "在腾讯公司，Jack是一名主管，主要负责天美旗下游戏。\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "19048b72-9a39-427d-9a71-7cf1dae7f1e7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'\\nwe will use the add() function to add text data with metadata and unique IDs. After that, Chroma will automatically download the all-MiniLM-L6-v2 model to convert the text into embeddings and store it in the \"Students\" collection.\\n'"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"\n",
    "we will use the add() function to add text data with metadata and unique IDs. After that, Chroma will automatically download the all-MiniLM-L6-v2 model to convert the text into embeddings and store it in the \"Students\" collection.\n",
    "\"\"\"\n",
    "\n",
    "# collection.add(\n",
    "#     documents = [student_info, club_info, university_info],\n",
    "#     metadatas = [{\"source\": \"student info\"},{\"source\": \"club info\"},{'source':'university info'}],\n",
    "#     ids = [\"id1\", \"id2\", \"id3\"]\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f6237d13-9ffe-4b50-b508-6f14e722211c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "支持的参数: ['self', 'api_key', 'model_name', 'organization_id', 'api_base', 'api_type', 'api_version', 'deployment_id', 'default_headers', 'dimensions', 'api_key_env_var']\n"
     ]
    }
   ],
   "source": [
    "import inspect\n",
    "sig = inspect.signature(embedding_functions.OpenAIEmbeddingFunction.__init__)\n",
    "print(\"支持的参数:\", list(sig.parameters.keys()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "19e8ebda-5110-4693-b03a-762117161b8f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[array([ 0.02474505, -0.06331136,  0.03730075, ...,  0.06820709,\n",
      "       -0.0294576 ,  0.06554275], shape=(1024,), dtype=float32), array([ 0.00679254, -0.05711098,  0.02125649, ..., -0.01073252,\n",
      "        0.0113209 ,  0.0479053 ], shape=(1024,), dtype=float32), array([ 0.02382195, -0.04835792,  0.04501505, ...,  0.04134764,\n",
      "        0.01528629,  0.05536819], shape=(1024,), dtype=float32)]\n",
      "1024\n",
      "1024\n",
      "1024\n"
     ]
    }
   ],
   "source": [
    "# ref https://help.aliyun.com/zh/model-studio/embedding-interfaces-compatible-with-openai?spm=5176.21213303.J_ZGek9Blx07Hclc3Ddt9dg.41.70b52f3dOmiBPs&scm=20140722.S_help@@%E6%96%87%E6%A1%A3@@2846066._.ID_help@@%E6%96%87%E6%A1%A3@@2846066-RL_OpenAIEmbeddingFunction-LOC_2024SPAllResult-OR_ser-PAR1_2150458217563022063747568eb88f-V_4-PAR3_o-RE_new8-P0_0-P1_0\n",
    "openai_ef = embedding_functions.OpenAIEmbeddingFunction(\n",
    "                model_name=\"text-embedding-v4\",\n",
    "                api_key=os.getenv(\"DASHSCOPE_API_KEY\"),\n",
    "                api_base=\"https://dashscope.aliyuncs.com/compatible-mode/v1\",\n",
    "                dimensions=1024,\n",
    "            )\n",
    "employers_embeddings = openai_ef([employer1_info, employer2_info, employer3_info])\n",
    "\n",
    "print(employers_embeddings)\n",
    "for employers_embedding in employers_embeddings:\n",
    "    print(len(employers_embedding))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "712352a5-2883-45b7-9015-b6e42885f7b0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "支持的参数: ['query_embeddings', 'query_texts', 'query_images', 'query_uris', 'ids', 'n_results', 'where', 'where_document', 'include']\n",
      "支持的参数: ['ids', 'embeddings', 'metadatas', 'documents', 'images', 'uris']\n"
     ]
    }
   ],
   "source": [
    "# https://docs.trychroma.com/docs/collections/update-data\n",
    "\n",
    "sig = inspect.signature(collection.query)\n",
    "print(\"支持的参数:\", list(sig.parameters.keys()))\n",
    "sig = inspect.signature(collection.add)\n",
    "print(\"支持的参数:\", list(sig.parameters.keys()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "37c40ce5-28dd-43d8-aadc-173229d52da5",
   "metadata": {},
   "outputs": [],
   "source": [
    "collection.add(\n",
    "    embeddings = employers_embeddings,\n",
    "    documents = [employer1_info, employer2_info, employer3_info],\n",
    "    metadatas = [{\"source\": \"employer1 info\"},{\"source\": \"employer2 info\"},{'source':'employer3 info'}],\n",
    "    ids = [\"id1\", \"id2\", \"id3\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "78eaf1c5-c75e-41d5-b338-7ce1592b1601",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'ids': [['id2', 'id1']], 'embeddings': None, 'documents': [['\\nTom是华为公司终端的一名销售，负责问界汽车的营销工作。\\n', '\\n我叫shinerio，来自北京，负责小米生态产品营销工作。\\n']], 'uris': None, 'included': ['metadatas', 'documents', 'distances'], 'data': None, 'metadatas': [[{'source': 'employer2 info'}, {'source': 'employer1 info'}]], 'distances': [[1.3764472007751465, 1.4127511978149414]]}\n"
     ]
    }
   ],
   "source": [
    "results = collection.query(\n",
    "    query_texts=[\"我想要买辆汽车\"],\n",
    "    n_results=2\n",
    ")\n",
    "print(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "bec953be-0dd5-43b4-9d79-c712e24aab95",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'ids': [['id3', 'id1']], 'embeddings': None, 'documents': [['\\n在腾讯公司，Jack是一名主管，主要负责天美旗下游戏。\\n', '\\n我叫shinerio，来自北京，负责小米生态产品营销工作。\\n']], 'uris': None, 'included': ['metadatas', 'documents', 'distances'], 'data': None, 'metadatas': [[{'source': 'employer3 info'}, {'source': 'employer1 info'}]], 'distances': [[1.4380024671554565, 1.4919731616973877]]}\n"
     ]
    }
   ],
   "source": [
    "results = collection.query(\n",
    "    query_texts=[\"什么游戏最好玩\"],\n",
    "    n_results=2\n",
    ")\n",
    "print(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "e7cbf1c0-d76e-40a6-892a-5ed33d090a36",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'ids': [['id1', 'id3']], 'embeddings': None, 'documents': [['\\n我叫shinerio，来自北京，负责小米生态产品营销工作。\\n', '\\n在腾讯公司，Jack是一名主管，主要负责天美旗下游戏。\\n']], 'uris': None, 'included': ['metadatas', 'documents', 'distances'], 'data': None, 'metadatas': [[{'source': 'employer1 info'}, {'source': 'employer3 info'}]], 'distances': [[1.321440577507019, 1.5759596824645996]]}\n"
     ]
    }
   ],
   "source": [
    "results = collection.query(\n",
    "    query_texts=[\"首都有什么好玩的地方\"],\n",
    "    n_results=2\n",
    ")\n",
    "print(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "38a94f42-3714-4d4e-a8bf-29cb88159bf3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'ids': [['id3', 'id1']], 'embeddings': None, 'documents': [['\\n在腾讯公司，Jack是一名主管，主要负责天美旗下游戏。\\n', '\\n我叫shinerio，来自北京，负责小米生态产品营销工作。\\n']], 'uris': None, 'included': ['metadatas', 'documents', 'distances'], 'data': None, 'metadatas': [[{'source': 'employer3 info'}, {'source': 'employer1 info'}]], 'distances': [[0.45964381098747253, 1.1920950412750244]]}\n"
     ]
    }
   ],
   "source": [
    "results = collection.query(\n",
    "    query_texts=[\"jack在腾讯公司主要负责什么业务\"],\n",
    "    n_results=2\n",
    ")\n",
    "print(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "3482d7c7-eec4-4d81-be4b-b5458156507a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'ids': [['id3', 'id2']], 'embeddings': None, 'documents': [['\\n在腾讯公司，Jack是一名主管，主要负责天美旗下游戏。\\n', '\\nTom是华为公司终端的一名销售，负责问界汽车的营销工作。\\n']], 'uris': None, 'included': ['metadatas', 'documents', 'distances'], 'data': None, 'metadatas': [[{'source': 'employer3 info'}, {'source': 'employer2 info'}]], 'distances': [[0.024908192455768585, 1.0855298042297363]]}\n"
     ]
    }
   ],
   "source": [
    "results = collection.query(\n",
    "    query_texts=[\"在腾讯公司，Jack是一名主管，主要负责天美旗下游戏。\"],\n",
    "    n_results=2\n",
    ")\n",
    "print(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b31950f0-3a08-49a8-aadd-1459097e6532",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python313",
   "language": "python",
   "name": "python313"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
